from typing import Any


class StopRecursive(Exception):
    """Exception to stop recursion in extraction."""

    def __init__(self, message: str):
        self.message = message

    def __str__(self) -> str:
        return self.message


class InvalidSchema(Exception):
    """Exception for invalid schema definitions."""

    pass


class ArraySchema:
    """Class to represent an array schema."""

    def __init__(self, schema: dict[Any, Any], name: str | None = None):
        self.schema = schema
        self.name = name

    def items(self) -> dict[Any, Any]:
        return self.schema.items()


class FieldSchema:
    """Class to represent a field schema."""

    def __init__(self, name: str, schema: dict[Any, Any]):
        self.name = name
        self.schema = schema

    def items(self) -> dict[Any, Any]:
        return self.schema.items()


def expand_dict(result: dict[Any, Any], new_data: dict[Any, Any]) -> None:
    for key, value in new_data.items():
        result[key] = value


def _extract_dict_with_schema(data: dict[Any, Any], schema: Any) -> dict[Any, Any]:
    result = {}
    for key, subschema in schema.items():
        try:
            extracted = None
            try:
                extracted = extract(data[key], subschema)
            except StopRecursive:
                result[
                    subschema.name if isinstance(subschema, FieldSchema) else key
                ] = data[key]
                continue
            if isinstance(subschema, FieldSchema):
                result[subschema.name] = extracted
                continue
            elif isinstance(subschema, ArraySchema) and subschema.name:
                result[subschema.name] = extracted
                continue
            expand_dict(result, extracted)
        except KeyError as e:
            print(f"KeyError: {e}")
            continue
        except IndexError as e:
            print(f"IndexError: {e}")
            continue
    return result


def extract(data: Any, schema: Any) -> Any:
    """
    Extracts data from a nested structure based on a provided schema.

    Args:
        data: The input data which can be a dictionary or a list.
        schema: A dictionary defining the structure to extract.

    Returns:
        The extracted data structured according to the schema.
    """
    if isinstance(data, list) and isinstance(schema, ArraySchema):
        return [
            _extract_dict_with_schema(
                item,
                schema.schema,
            )
            for item in data
        ]

    if not isinstance(data, dict) and not isinstance(data, list):
        raise StopRecursive("Data is not a dictionary")

    result = _extract_dict_with_schema(data, schema)

    return result


def test_extract():
    data = {
        "name": "John",
        "age": 30,
        "address": {"street": "123 Main St", "city": "New York", "zip": "10001"},
        "phones": [
            {"type": "home", "number": "212-555-1234"},
            {"type": "work", "number": "646-555-5678"},
        ],
    }

    schema = {
        "name": FieldSchema("full_name", {}),
        "address": {"city": {}, "zip": FieldSchema("zip", {})},
        "phones": {
            0: {"number": {}},
        },
    }

    extracted = extract(data, schema)
    expected = {
        "full_name": "John",
        "city": "New York",
        "zip": "10001",
        "number": "212-555-1234"
    }
    print(extracted)

    assert extracted == expected, f"Expected {expected}, but got {extracted}"


if __name__ == "__main__":
    test_extract()
